import threading
import time
import paramiko
import socket
from io import BytesIO
from urllib.parse import urlparse
from typing import Optional, Dict
from scapy.all import Raw
from scapy.layers.inet import TCP
from scapy.utils import rdpcap
from base.base_page import BasePage
from pages.http_pages import HttpRequestPage
from config.settings import get_config


class NetworkCaptureException(Exception):
    """自定义网络捕获异常"""
    pass


class Ssh_O_Http_Req_Header(BasePage):
    def __init__(self, driver, device_type, logger=None):
        super().__init__(driver, device_type, logger=logger)
        self.driver = driver
        self.config = get_config(device_type)
        self.http_client = HttpRequestPage(driver, device_type, logger=logger)
        self.capture_lock = threading.Lock()
        self.capture_buffer = BytesIO()
        self.capture_active = False

    def capture_and_request(
            self,
            ssh_config: Dict,
            http_config: Dict,
            capture_timeout: int = 2
    ) -> Optional[str]:
        """
        整合抓包与HTTP请求的核心方法

        :param ssh_config: SSH配置字典
            {
                "ip": "192.168.1.100",
                "port": 22,
                "username": "admin",
                "password": "secret"
            }
        :param http_config: HTTP请求配置
            {
                "url": "http://example.com",
                "method": "GET",
                "headers": {"X-Header": "value"},
                "body": "data",
                "apikey": "your-api-key",
                "target_header": "Authorization"
            }
        :param capture_timeout: 最大抓包等待时间(秒)
        """
        try:
            # 验证配置
            self._validate_config(ssh_config, http_config)

            # 解析目标地址
            target_ip = ssh_config['ip']

            # 清空缓冲区
            self.capture_buffer = BytesIO()

            # 启动SSH抓包线程
            ssh_thread = threading.Thread(
                target=self._ssh_capture_worker,
                args=(ssh_config, target_ip, capture_timeout)
            )
            ssh_thread.start()

            # 等待抓包初始化
            time.sleep(1)  # 增加等待时间确保抓包启动

            # 发送HTTP请求
            http_status, http_response = self._execute_http_request(http_config)
            self.logger.info(f"HTTP请求完成，状态码: {http_status}")

            # 等待抓包结束
            self.capture_active = False
            ssh_thread.join(timeout=capture_timeout + 5)

            # 解析捕获数据
            return self.parse_captured_data(http_config['target_header'])

        except Exception as e:
            self.logger.error(f"操作失败: {str(e)}", exc_info=True)
            raise NetworkCaptureException(f"网络捕获失败: {str(e)}")

    def _ssh_capture_worker(self, config: Dict, target_ip: str, timeout: int):
        """免sudo抓包工作线程"""
        client = None
        try:
            client = self._create_ssh_client(config)
            self.logger.info("SSH连接已建立，启动抓包...")

            # 使用完整路径避免PATH问题
            cmd = (
                f"/usr/sbin/tcpdump -i any -s 0 -U -w - "
                f"'(host {target_ip} and tcp port 5000)' "
                f"-G {timeout} -W 1"  # 替代timeout命令
            )
            self.logger.info(f"执行命令: {cmd}")

            chan = client.get_transport().open_session()
            chan.exec_command(cmd)

            # 设置非阻塞模式
            chan.setblocking(0)

            # 直接开始捕获数据
            start_time = time.time()
            while time.time() - start_time < timeout + 5:
                if chan.recv_ready():
                    data = chan.recv(4096)
                    self.capture_buffer.write(data)
                if chan.recv_stderr_ready():
                    error = chan.recv_stderr(4096)
                    self.logger.debug(f"tcpdump错误输出: {error.decode()}")
                time.sleep(0.01)

            # 读取最终数据
            while chan.recv_ready():
                data = chan.recv(4096)
                self.capture_buffer.write(data)

        except Exception as e:
            self.logger.error(f"抓包异常: {str(e)}")
        finally:
            if client:
                client.close()
            self.logger.info("抓包线程安全退出")

    def _clean_ssh_client(self, client: Optional[paramiko.SSHClient]):
        """安全清理SSH连接"""
        try:
            if client:
                # 发送优雅关闭信号
                transport = client.get_transport()
                if transport and transport.is_active():
                    transport.close()
                # 双重保障关闭
                client.close()
        except Exception as e:
            self.logger.warning(f"清理SSH连接时发生异常: {str(e)}")

    def _create_ssh_client(self, config: Dict) -> paramiko.SSHClient:
        """创建SSH连接（增强版错误诊断）"""
        max_retries = 5
        retry_delay = 1  # 基础等待时间（秒）
        client = None
        host = config['ip']
        port = config.get('port', 22)

        # 预检网络连通性
        def check_port_access():
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(3)
            try:
                result = sock.connect_ex((host, port))
                if result == 0:
                    self.logger.info(f"端口 {port} 可达性检查通过")
                else:
                    raise NetworkCaptureException(
                        f"端口 {port} 不可达 (错误码: {result})，请检查：\n"
                        f"1. 目标主机 {host} 是否在线\n"
                        f"2. 防火墙是否开放端口 {port}\n"
                        f"3. SSH服务是否运行 (sudo systemctl status sshd)"
                    )
            finally:
                sock.close()

        for attempt in range(max_retries):
            try:
                # 执行网络预检
                if attempt == 0:
                    check_port_access()

                client = paramiko.SSHClient()
                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                # 增强连接参数配置
                connect_params = {
                    'hostname': host,
                    'port': port,
                    'username': config['username'],
                    'password': config['password'],
                    'timeout': 5,
                    'banner_timeout': 10,
                    'auth_timeout': 8,
                    'look_for_keys': False,
                    'allow_agent': False
                }

                # 诊断模式：首次尝试显示详细日志
                if attempt == 0:
                    paramiko.common.logging.basicConfig(level=paramiko.common.DEBUG)

                client.connect(**connect_params)

                # 恢复日志级别
                paramiko.common.logging.basicConfig(level=paramiko.common.WARNING)

                # 权限验证流程优化
                self._validate_ssh_permission(client)
                return client

            except (socket.timeout, TimeoutError) as e:
                error_type = "TCP连接超时"
                self.logger.warning(f"{error_type}（尝试 {attempt + 1}/{max_retries}）")
                self._clean_ssh_client(client)
                if attempt == max_retries - 1:
                    raise NetworkCaptureException(
                        f"连续连接超时，请检查：\n"
                        f"1. 网络延迟是否过高\n"
                        f"2. 目标主机 {host} 的SSH服务负载\n"
                        f"3. 中间网络设备（如防火墙）的会话限制"
                    ) from e

            except paramiko.AuthenticationException as e:
                raise NetworkCaptureException(
                    f"认证失败，请确认：\n"
                    f"1. 用户名 '{config['username']}' 是否正确\n"
                    f"2. 密码是否有效\n"
                    f"3. 账户是否被锁定"
                ) from e

            except paramiko.SSHException as e:
                error_msg = str(e).lower()
                diagnosis = ""

                if "connection refused" in error_msg:
                    diagnosis = "SSH服务未运行或端口错误"
                elif "no existing session" in error_msg:
                    diagnosis = "SSH会话初始化失败"
                elif "error reading ssh protocol banner" in error_msg:
                    diagnosis = "网络不稳定或SSH服务配置异常"

                self.logger.error(f"SSH协议错误: {diagnosis}（{error_msg}）")
                if attempt == max_retries - 1:
                    raise NetworkCaptureException(
                        f"SSH协议层错误: {diagnosis}\n"
                        f"建议检查：\n"
                        f"1. /etc/ssh/sshd_config 配置\n"
                        f"2. 系统日志 /var/log/auth.log\n"
                        f"3. 使用telnet {host} {port} 验证协议握手"
                    ) from e

            except socket.error as e:
                error_code = e.errno if hasattr(e, 'errno') else 'N/A'
                error_msg_map = {
                    113: "目标网络不可达，检查路由表",
                    111: "连接被拒绝，确认服务状态",
                    101: "网络不可达，检查VPN连接",
                    99: "地址绑定错误，尝试重启网络服务"
                }
                diagnosis = error_msg_map.get(error_code, "检查网络基础设施")

                raise NetworkCaptureException(
                    f"底层网络错误 [{error_code}]: {diagnosis}\n"
                    f"详细原因: {str(e)}"
                ) from e

            except Exception as e:
                self.logger.error("未知连接错误", exc_info=True)
                raise NetworkCaptureException(
                    f"未预期的错误类型: {type(e).__name__}\n"
                    f"建议收集以下信息进行诊断：\n"
                    f"1. 从客户端执行: nc -zv {host} {port}\n"
                    f"2. 从客户端执行: traceroute {host}\n"
                    f"3. 目标主机 {host} 的系统日志"
                ) from e

            # 指数退避策略
            time.sleep(retry_delay * (2  ** attempt))

        raise NetworkCaptureException("连接流程异常终止")

    def _validate_ssh_permission(self, client: paramiko.SSHClient):
        """增强权限验证流程"""
        try:
            # 执行多层级权限检查
            commands = [
                ("基础Shell访问", "echo $SHELL"),  # 验证基本shell访问
                ("Sudo权限", "sudo -n true"),  # 验证无密码sudo
                ("抓包工具", "which tcpdump"),  # 验证工具安装
                ("网络接口", "ip link show")  # 验证接口查看权限
            ]

            for check_name, cmd in commands:
                stdin, stdout, stderr = client.exec_command(cmd, timeout=5)
                exit_code = stdout.channel.recv_exit_status()

                if exit_code != 0:
                    error_output = stderr.read().decode().strip()
                    raise PermissionError(
                        f"权限检查失败 [{check_name}]:\n"
                        f"命令: {cmd}\n"
                        f"错误: {error_output or '无输出'}"
                    )

        except paramiko.SSHException as e:
            raise NetworkCaptureException("权限验证期间连接中断") from e

    def _execute_http_request(self, config: Dict) -> tuple:
        """执行HTTP请求"""
        return self.http_client.http_request(
            direction="请求方",
            url=config['url'],
            method=config.get('method', 'GET'),
            headers=config.get('headers', {}),
            body=config.get('body'),
            apikey=config['apikey'],
            expected_status=200,
            timeout=config.get('timeout', 5)
        )

    def parse_captured_data(self, target_header: str) -> Optional[str]:
        try:
            raw_data = self.capture_buffer.getvalue()
            if len(raw_data) < 24:  # pcap文件头最小长度
                self.logger.error("捕获数据过短")
                return None

            # 检查magic number
            magic_number = raw_data[:4]
            if magic_number not in {b"\xa1\xb2\xc3\xd4", b"\xd4\xc3\xb2\xa1"}:
                self.logger.error(f"无效的pcap文件头: {magic_number.hex()}")
                return None

            # 保存调试文件
            # with open("debug_capture.pcap", "wb") as f:
            #     f.write(raw_data)

            # 使用Scapy解析
            packets = rdpcap(BytesIO(raw_data))
            self.logger.info(f"成功解析到 {len(packets)} 个数据包")

            for pkt in packets:
                if self._is_http_request(pkt):
                    headers = self._extract_http_headers(pkt)
                    self.logger.debug(f"解析到HTTP头: {headers}")
                    if target_header in headers:
                        self.logger.info(f"找到目标头: {target_header}={headers[target_header]}")
                        return headers[target_header]
            return None
        except Exception as e:
            self.logger.error(f"数据解析失败: {str(e)}", exc_info=True)
            return None

    def _is_http_request(self, packet) -> bool:
        """判断是否为HTTP请求包"""
        try:
            return (
                    packet.haslayer(TCP) and
                    packet[TCP].dport == 80 and
                    packet.haslayer(Raw) and
                    b'HTTP' in bytes(packet[Raw].load)
            )
        except:
            return False

    def _extract_http_headers(self, packet) -> Dict:
        """从数据包提取HTTP头"""
        try:
            payload = packet[Raw].load.decode('utf-8', 'ignore')
            header_block = payload.split('\r\n\r\n', 1)[0]
            headers = {}
            for line in header_block.split('\r\n')[1:]:
                if ': ' in line:
                    key, value = line.split(': ', 1)
                    headers[key.strip()] = value.strip()
            return headers
        except Exception as e:
            self.logger.debug(f"头解析异常: {str(e)}")
            return {}

    def resolve_target(self, url: str) -> str:
        """解析目标IP地址"""
        parsed = urlparse(url)
        if not parsed.hostname:
            raise ValueError("无效URL")
        try:
            return socket.gethostbyname(parsed.hostname)
        except socket.gaierror:
            raise NetworkCaptureException("DNS解析失败")

    def _validate_config(self, ssh_cfg: Dict, http_cfg: Dict):
        """验证配置完整性"""
        required_ssh = ['ip', 'username', 'password']
        required_http = ['url', 'apikey', 'target_header']

        missing_ssh = [k for k in required_ssh if k not in ssh_cfg]
        missing_http = [k for k in required_http if k not in http_cfg]

        if missing_ssh:
            raise ValueError(f"SSH配置缺少参数: {missing_ssh}")
        if missing_http:
            raise ValueError(f"HTTP配置缺少参数: {missing_http}")

